from pymilvus import MilvusClient
from dotenv import load_dotenv
import os
from vector_database import chunk, embed

load_dotenv()

def get_client() -> MilvusClient:
    client = MilvusClient(
        uri=os.getenv("MILVUS_URI"),
        user=os.getenv("MILVUS_USER"),
        password=os.getenv("MILVUS_PASSWORD")
    )
    if not client.has_collection(collection_name=os.getenv("MILVUS_COLLECTION")):
        client.create_collection(
            collection_name=os.getenv("MILVUS_COLLECTION"),
            dimension=768
        )
    
    return client


def get_records(filename):
    data = []
    chunks = chunk.get_chunks(filename)
    total = len(chunks)
    for i, c in enumerate(chunks):
        text = c.page_content
        metadata = c.metadata
        print(f"({i + 1}/{total}) ", end="")
        vector = embed.embed(text)

        record = {
            "id": i,
            "vector": vector,
            "text": text,
            "chapters": metadata,
        }
        data.append(record)

    return data


def insert_db(records: list[dict]):
    client = get_client()
    data = records
    res = client.insert(
        collection_name=os.getenv("MILVUS_COLLECTION"),
        data=data
    )
    return res


def query_db(question: str, limit: int = 5) -> list[str]:
    client = get_client()
    question_embeddings = embed.embed(question)

    res = client.search(
        collection_name=os.getenv("MILVUS_COLLECTION"),
        data = [question_embeddings],
        limit=limit,
        output_fields=["text", "chapters"]
    )

    return res[0]
